{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple/\n",
      "Requirement already satisfied: dashscope in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (1.18.0)\n",
      "Requirement already satisfied: vanna[chromadb,mysql] in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (0.5.0)\n",
      "Requirement already satisfied: requests in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (2.31.0)\n",
      "Requirement already satisfied: tabulate in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (0.9.0)\n",
      "Requirement already satisfied: plotly in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (5.21.0)\n",
      "Requirement already satisfied: pandas in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (2.2.2)\n",
      "Requirement already satisfied: sqlparse in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (0.5.0)\n",
      "Requirement already satisfied: kaleido in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (0.2.1)\n",
      "Requirement already satisfied: flask in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (3.0.3)\n",
      "Requirement already satisfied: flask-sock in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (0.7.0)\n",
      "Requirement already satisfied: sqlalchemy in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (2.0.29)\n",
      "Requirement already satisfied: chromadb in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (0.5.0)\n",
      "Requirement already satisfied: PyMySQL in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from vanna[chromadb,mysql]) (1.1.0)\n",
      "Requirement already satisfied: aiohttp in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from dashscope) (3.9.5)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from aiohttp->dashscope) (4.0.3)\n",
      "Requirement already satisfied: build>=1.0.3 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.2.1)\n",
      "Requirement already satisfied: pydantic>=1.9 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (2.7.1)\n",
      "Requirement already satisfied: chroma-hnswlib==0.7.3 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.7.3)\n",
      "Requirement already satisfied: fastapi>=0.95.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.110.3)\n",
      "Requirement already satisfied: uvicorn>=0.18.3 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (0.29.0)\n",
      "Requirement already satisfied: numpy>=1.22.5 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.26.4)\n",
      "Requirement already satisfied: posthog>=2.4.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (3.5.0)\n",
      "Requirement already satisfied: typing-extensions>=4.5.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (4.11.0)\n",
      "Requirement already satisfied: onnxruntime>=1.14.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.17.3)\n",
      "Requirement already satisfied: opentelemetry-api>=1.2.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-grpc>=1.2.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-fastapi>=0.41b0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-sdk>=1.2.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.24.0)\n",
      "Requirement already satisfied: tokenizers>=0.13.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.19.1)\n",
      "Requirement already satisfied: pypika>=0.48.9 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.48.9)\n",
      "Requirement already satisfied: tqdm>=4.65.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (4.66.2)\n",
      "Requirement already satisfied: overrides>=7.3.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (7.7.0)\n",
      "Requirement already satisfied: importlib-resources in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (6.4.0)\n",
      "Requirement already satisfied: grpcio>=1.58.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (1.62.2)\n",
      "Requirement already satisfied: bcrypt>=4.0.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (4.1.2)\n",
      "Requirement already satisfied: typer>=0.9.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (0.12.3)\n",
      "Requirement already satisfied: kubernetes>=28.1.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (29.0.0)\n",
      "Requirement already satisfied: tenacity>=8.2.3 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (8.2.3)\n",
      "Requirement already satisfied: PyYAML>=6.0.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (6.0.1)\n",
      "Requirement already satisfied: mmh3>=4.0.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (4.1.0)\n",
      "Requirement already satisfied: orjson>=3.9.12 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from chromadb->vanna[chromadb,mysql]) (3.10.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from requests->vanna[chromadb,mysql]) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from requests->vanna[chromadb,mysql]) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from requests->vanna[chromadb,mysql]) (2.2.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from requests->vanna[chromadb,mysql]) (2024.2.2)\n",
      "Requirement already satisfied: Werkzeug>=3.0.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask->vanna[chromadb,mysql]) (3.0.2)\n",
      "Requirement already satisfied: Jinja2>=3.1.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask->vanna[chromadb,mysql]) (3.1.3)\n",
      "Requirement already satisfied: itsdangerous>=2.1.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask->vanna[chromadb,mysql]) (2.2.0)\n",
      "Requirement already satisfied: click>=8.1.3 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask->vanna[chromadb,mysql]) (8.1.7)\n",
      "Requirement already satisfied: blinker>=1.6.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask->vanna[chromadb,mysql]) (1.8.1)\n",
      "Requirement already satisfied: simple-websocket>=0.5.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from flask-sock->vanna[chromadb,mysql]) (1.0.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pandas->vanna[chromadb,mysql]) (2.9.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pandas->vanna[chromadb,mysql]) (2024.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pandas->vanna[chromadb,mysql]) (2024.1)\n",
      "Requirement already satisfied: packaging in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from plotly->vanna[chromadb,mysql]) (24.0)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from sqlalchemy->vanna[chromadb,mysql]) (3.0.3)\n",
      "Requirement already satisfied: pyproject_hooks in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from build>=1.0.3->chromadb->vanna[chromadb,mysql]) (1.1.0)\n",
      "Requirement already satisfied: tomli>=1.1.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from build>=1.0.3->chromadb->vanna[chromadb,mysql]) (2.0.1)\n",
      "Requirement already satisfied: starlette<0.38.0,>=0.37.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from fastapi>=0.95.2->chromadb->vanna[chromadb,mysql]) (0.37.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from Jinja2>=3.1.2->flask->vanna[chromadb,mysql]) (2.1.5)\n",
      "Requirement already satisfied: six>=1.9.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (1.16.0)\n",
      "Requirement already satisfied: google-auth>=1.0.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (2.29.0)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (1.8.0)\n",
      "Requirement already satisfied: requests-oauthlib in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (2.0.0)\n",
      "Requirement already satisfied: oauthlib>=3.2.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (3.2.2)\n",
      "Requirement already satisfied: coloredlogs in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (24.3.25)\n",
      "Requirement already satisfied: protobuf in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (4.25.3)\n",
      "Requirement already satisfied: sympy in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (1.12)\n",
      "Requirement already satisfied: deprecated>=1.2.6 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-api>=1.2.0->chromadb->vanna[chromadb,mysql]) (1.2.14)\n",
      "Requirement already satisfied: importlib-metadata<=7.0,>=6.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-api>=1.2.0->chromadb->vanna[chromadb,mysql]) (7.0.0)\n",
      "Requirement already satisfied: googleapis-common-protos~=1.52 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb->vanna[chromadb,mysql]) (1.63.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-common==1.24.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb->vanna[chromadb,mysql]) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-proto==1.24.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb->vanna[chromadb,mysql]) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-asgi==0.45b0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation==0.45b0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-semantic-conventions==0.45b0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-util-http==0.45b0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (0.45b0)\n",
      "Requirement already satisfied: setuptools>=16.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (68.2.2)\n",
      "Requirement already satisfied: wrapt<2.0.0,>=1.0.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (1.16.0)\n",
      "Requirement already satisfied: asgiref~=3.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from opentelemetry-instrumentation-asgi==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb->vanna[chromadb,mysql]) (3.8.1)\n",
      "Requirement already satisfied: monotonic>=1.5 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from posthog>=2.4.0->chromadb->vanna[chromadb,mysql]) (1.6)\n",
      "Requirement already satisfied: backoff>=1.10.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from posthog>=2.4.0->chromadb->vanna[chromadb,mysql]) (2.2.1)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pydantic>=1.9->chromadb->vanna[chromadb,mysql]) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.18.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pydantic>=1.9->chromadb->vanna[chromadb,mysql]) (2.18.2)\n",
      "Requirement already satisfied: wsproto in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from simple-websocket>=0.5.1->flask-sock->vanna[chromadb,mysql]) (1.2.0)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from tokenizers>=0.13.2->chromadb->vanna[chromadb,mysql]) (0.22.2)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from typer>=0.9.0->chromadb->vanna[chromadb,mysql]) (1.5.4)\n",
      "Requirement already satisfied: rich>=10.11.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from typer>=0.9.0->chromadb->vanna[chromadb,mysql]) (13.7.1)\n",
      "Requirement already satisfied: h11>=0.8 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn>=0.18.3->uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (0.14.0)\n",
      "Requirement already satisfied: httptools>=0.5.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (0.6.1)\n",
      "Requirement already satisfied: python-dotenv>=0.13 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (1.0.1)\n",
      "Requirement already satisfied: uvloop!=0.15.0,!=0.15.1,>=0.14.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (0.19.0)\n",
      "Requirement already satisfied: watchfiles>=0.13 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (0.21.0)\n",
      "Requirement already satisfied: websockets>=10.4 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from uvicorn[standard]>=0.18.3->chromadb->vanna[chromadb,mysql]) (12.0)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (5.3.3)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (0.4.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (4.9)\n",
      "Requirement already satisfied: filelock in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb->vanna[chromadb,mysql]) (3.14.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from huggingface-hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb->vanna[chromadb,mysql]) (2024.3.1)\n",
      "Requirement already satisfied: zipp>=0.5 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from importlib-metadata<=7.0,>=6.0->opentelemetry-api>=1.2.0->chromadb->vanna[chromadb,mysql]) (3.17.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from rich>=10.11.0->typer>=0.9.0->chromadb->vanna[chromadb,mysql]) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from rich>=10.11.0->typer>=0.9.0->chromadb->vanna[chromadb,mysql]) (2.17.2)\n",
      "Requirement already satisfied: anyio<5,>=3.4.0 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from starlette<0.38.0,>=0.37.2->fastapi>=0.95.2->chromadb->vanna[chromadb,mysql]) (4.3.0)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from coloredlogs->onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (10.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from sympy->onnxruntime>=1.14.1->chromadb->vanna[chromadb,mysql]) (1.3.0)\n",
      "Requirement already satisfied: sniffio>=1.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from anyio<5,>=3.4.0->starlette<0.38.0,>=0.37.2->fastapi>=0.95.2->chromadb->vanna[chromadb,mysql]) (1.3.1)\n",
      "Requirement already satisfied: exceptiongroup>=1.0.2 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from anyio<5,>=3.4.0->starlette<0.38.0,>=0.37.2->fastapi>=0.95.2->chromadb->vanna[chromadb,mysql]) (1.2.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer>=0.9.0->chromadb->vanna[chromadb,mysql]) (0.1.2)\n",
      "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /home/owenliang/.conda/envs/vanna/lib/python3.10/site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=28.1.0->chromadb->vanna[chromadb,mysql]) (0.6.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install 'vanna[chromadb,mysql]' dashscope -i https://mirrors.aliyun.com/pypi/simple/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vanna.base import VannaBase\n",
    "from vanna.chromadb import ChromaDB_VectorStore\n",
    "from dashscope import Generation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 自定义LLM\n",
    "\n",
    "https://vanna.ai/docs/mysql-other-llm-chromadb/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random \n",
    "\n",
    "DEBUG_INFO=None\n",
    "\n",
    "class QwenLLM(VannaBase):\n",
    "  def __init__(self,config=None):\n",
    "    self.model=config['model']\n",
    "    self.api_key=config['api_key']\n",
    "  \n",
    "  def system_message(self,message: str):\n",
    "    return {'role':'system','content':message}\n",
    "\n",
    "  def user_message(self, message: str):\n",
    "    return {'role':'user','content':message}\n",
    "\n",
    "  def assistant_message(self, message: str):\n",
    "    return {'role':'assistant','content':message}\n",
    "  \n",
    "  def submit_prompt(self,prompt,**kwargs):\n",
    "    resp=Generation.call(\n",
    "      model=self.model,\n",
    "      messages=prompt,\n",
    "      seed=random.randint(1, 10000),\n",
    "      result_format='message',\n",
    "      api_key=self.api_key)\n",
    "    answer=resp.output.choices[0].message.content\n",
    "    global DEBUG_INFO\n",
    "    DEBUG_INFO=(prompt,answer)\n",
    "    return answer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vanna客户端"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyVanna(ChromaDB_VectorStore,QwenLLM):\n",
    "    def __init__(self, config=None):\n",
    "        ChromaDB_VectorStore.__init__(self,config=config)\n",
    "        QwenLLM.__init__(self,config=config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "vn=MyVanna({'api_key':'sk-d8d15e769aa549129467e605c1f2905e','model':'qwen-max'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MySQL服务端"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "sudo docker run -d --name mysql-vanna -p 3306:3306 -e MYSQL_ROOT_PASSWORD=123456 mysql:8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "mysql --protocol=tcp -hlocalhost -P3306 -uroot -p123456"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "CREATE TABLE IF NOT EXISTS user (\n",
    "        id INT PRIMARY KEY COMMENT '用户ID' ,\n",
    "        name VARCHAR(100) COMMENT '姓名',\n",
    "        age INT COMMENT '年龄'\n",
    "    ) COMMENT '用户信息表';\n",
    "    \n",
    "insert into user values(1,'小鱼儿',34),(2,'小悲剧',36);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 构造向量库"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "vn.connect_to_mysql(host='localhost',dbname='vanna',user='root',password='root',port=3306)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "DDL='''CREATE TABLE IF NOT EXISTS user (\n",
    "        id INT PRIMARY KEY COMMENT '用户ID' ,\n",
    "        name VARCHAR(100) COMMENT '姓名',\n",
    "        age INT COMMENT '年龄'\n",
    "    ) COMMENT '用户信息表';\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding ddl: CREATE TABLE IF NOT EXISTS user (\n",
      "        id INT PRIMARY KEY COMMENT '用户ID' ,\n",
      "        name VARCHAR(100) COMMENT '姓名',\n",
      "        age INT COMMENT '年龄'\n",
      "    ) COMMENT '用户信息表';\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\101\\.cache\\chroma\\onnx_models\\all-MiniLM-L6-v2\\onnx.tar.gz: 100%|██████████| 79.3M/79.3M [07:11<00:00, 193kiB/s] \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'ab0ac208-2f5e-50b0-9177-423427220940-ddl'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 存储DDL到向量库\n",
    "vn.train(ddl=DDL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding documentation....\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'8fc54ebe-8bb3-5fb7-88a6-5c98d817ed07-doc'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 存储document到向量库\n",
    "vn.train(documentation='\"福报\"是指age>=35岁，也就是可以向社会输送的人才')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question generated with sql: Who are the users aged between 10 and 20? \n",
      "Adding SQL...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'04a88b26-6984-5521-b897-73798ce0001f-sql'"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 存储SQL到向量库\n",
    "'''\n",
    "1，通过LLM根据SQL构造一个question\n",
    "2，按question-SQL的JSON入库\n",
    "            {\n",
    "                \"question\": question,\n",
    "                \"sql\": sql,\n",
    "            }\n",
    "'''\n",
    "\n",
    "vn.train(sql='select name from user where age between 10 and 20')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROMPT: The user will give you SQL and you will try to guess what the business question this query is answering. Return just the question without any additional explanation. Do not reference the table name in the question.\n",
      "ANSWER: Who are the users aged between 10 and 20?\n"
     ]
    }
   ],
   "source": [
    "Q,A=DEBUG_INFO\n",
    "print('PROMPT:',Q[0]['content'])\n",
    "print('ANSWER:',A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'e5102160-2dbf-5300-98f5-24d762a12b59-sql'"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 存储question-SQL到向量库\n",
    "'''\n",
    "按question-SQL的JSON入库\n",
    "            {\n",
    "                \"question\": question,\n",
    "                \"sql\": sql,\n",
    "            }\n",
    "'''\n",
    "vn.train(question='小鱼儿的年龄',sql='select age from user where name=\"小鱼儿\"')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>question</th>\n",
       "      <th>content</th>\n",
       "      <th>training_data_type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>04a88b26-6984-5521-b897-73798ce0001f-sql</td>\n",
       "      <td>Who are the users aged between 10 and 20?</td>\n",
       "      <td>select name from user where age between 10 and 20</td>\n",
       "      <td>sql</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>e5102160-2dbf-5300-98f5-24d762a12b59-sql</td>\n",
       "      <td>小鱼儿的年龄</td>\n",
       "      <td>select age from user where name=\"小鱼儿\"</td>\n",
       "      <td>sql</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>ab0ac208-2f5e-50b0-9177-423427220940-ddl</td>\n",
       "      <td>None</td>\n",
       "      <td>CREATE TABLE IF NOT EXISTS user (\\n        id ...</td>\n",
       "      <td>ddl</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>8fc54ebe-8bb3-5fb7-88a6-5c98d817ed07-doc</td>\n",
       "      <td>None</td>\n",
       "      <td>\"福报\"是指age&gt;=35岁，也就是可以向社会输送的人才</td>\n",
       "      <td>documentation</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                         id  \\\n",
       "0  04a88b26-6984-5521-b897-73798ce0001f-sql   \n",
       "1  e5102160-2dbf-5300-98f5-24d762a12b59-sql   \n",
       "0  ab0ac208-2f5e-50b0-9177-423427220940-ddl   \n",
       "0  8fc54ebe-8bb3-5fb7-88a6-5c98d817ed07-doc   \n",
       "\n",
       "                                    question  \\\n",
       "0  Who are the users aged between 10 and 20?   \n",
       "1                                     小鱼儿的年龄   \n",
       "0                                       None   \n",
       "0                                       None   \n",
       "\n",
       "                                             content training_data_type  \n",
       "0  select name from user where age between 10 and 20                sql  \n",
       "1              select age from user where name=\"小鱼儿\"                sql  \n",
       "0  CREATE TABLE IF NOT EXISTS user (\\n        id ...                ddl  \n",
       "0                       \"福报\"是指age>=35岁，也就是可以向社会输送的人才      documentation  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 检查所有入库的知识\n",
    "vn.get_training_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 开始查询"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '用户的平均年龄'}]\n",
      "LLM Response: select avg(age) from user\n",
      "SQL: select avg(age) from user\n",
      "PROMPT: You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \n",
      "===Tables \n",
      "CREATE TABLE IF NOT EXISTS user (\n",
      "        id INT PRIMARY KEY COMMENT '用户ID' ,\n",
      "        name VARCHAR(100) COMMENT '姓名',\n",
      "        age INT COMMENT '年龄'\n",
      "    ) COMMENT '用户信息表';\n",
      "\n",
      "\n",
      "\n",
      "===Additional Context \n",
      "\n",
      "\"福报\"是指age>=35岁，也就是可以向社会输送的人才\n",
      "\n",
      "===Response Guidelines \n",
      "1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \n",
      "2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \n",
      "3. If the provided context is insufficient, please explain why it can't be generated. \n",
      "4. Please use the most relevant table(s). \n",
      "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n",
      "6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \n",
      "\n",
      "ANSWER: select avg(age) from user\n"
     ]
    }
   ],
   "source": [
    "# 基本使用\n",
    "result=vn.generate_sql('用户的平均年龄')\n",
    "print('SQL:',result)\n",
    "\n",
    "Q,A=DEBUG_INFO\n",
    "print('PROMPT:',Q[0]['content'])\n",
    "print('ANSWER:',A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '用户的平均年龄'}]\n",
      "LLM Response: select avg(age) from user\n",
      "select avg(age) from user\n",
      "Couldn't run sql:  (1146, \"Table 'vanna.user' doesn't exist\")\n"
     ]
    }
   ],
   "source": [
    "vn.ask('用户的平均年龄')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '打算给一批员工送福报，把他们的名字过滤出来'}]\n",
      "LLM Response: select name from user where age >= 35\n"
     ]
    }
   ],
   "source": [
    "result=vn.generate_sql('打算给一批员工送福报，把他们的名字过滤出来')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '打算给一批员工送福报，把他们的名字过滤出来'}]\n",
      "LLM Response: select name from user where age >= 35\n",
      "select name from user where age >= 35\n",
      "Couldn't run sql:  (1146, \"Table 'vanna.user' doesn't exist\")\n"
     ]
    }
   ],
   "source": [
    "vn.ask('打算给一批员工送福报，把他们的名字过滤出来')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PROMPT: You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \n",
      "===Tables \n",
      "CREATE TABLE IF NOT EXISTS user (\n",
      "        id INT PRIMARY KEY COMMENT '用户ID' ,\n",
      "        name VARCHAR(100) COMMENT '姓名',\n",
      "        age INT COMMENT '年龄'\n",
      "    ) COMMENT '用户信息表';\n",
      "\n",
      "\n",
      "\n",
      "===Additional Context \n",
      "\n",
      "\"福报\"是指age>=35岁，也就是可以向社会输送的人才\n",
      "\n",
      "===Response Guidelines \n",
      "1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \n",
      "2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \n",
      "3. If the provided context is insufficient, please explain why it can't be generated. \n",
      "4. Please use the most relevant table(s). \n",
      "5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \n",
      "6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \n",
      "\n",
      "ANSWER: select name from user where age >= 35\n"
     ]
    }
   ],
   "source": [
    "Q,A=DEBUG_INFO\n",
    "print('PROMPT:',Q[0]['content'])\n",
    "print('ANSWER:',A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '统计一下各年龄段的用户数量,年龄段是指0-10,10-20,20-30,30-40,40-50,50-60,60-70,70-80...左闭右开区间'}]\n",
      "LLM Response: ```sql\n",
      "SELECT \n",
      "    CASE \n",
      "        WHEN age >= 0 AND age < 10 THEN '0-10'\n",
      "        WHEN age >= 10 AND age < 20 THEN '10-20'\n",
      "        WHEN age >= 20 AND age < 30 THEN '20-30'\n",
      "        WHEN age >= 30 AND age < 40 THEN '30-40'\n",
      "        WHEN age >= 40 AND age < 50 THEN '40-50'\n",
      "        WHEN age >= 50 AND age < 60 THEN '50-60'\n",
      "        WHEN age >= 60 AND age < 70 THEN '60-70'\n",
      "        WHEN age >= 70 AND age < 80 THEN '70-80'\n",
      "        ELSE '80+'\n",
      "    END AS age_group,\n",
      "    COUNT(*) AS user_count\n",
      "FROM user\n",
      "GROUP BY age_group;\n",
      "```\n",
      "Extracted SQL: SELECT \n",
      "    CASE \n",
      "        WHEN age >= 0 AND age < 10 THEN '0-10'\n",
      "        WHEN age >= 10 AND age < 20 THEN '10-20'\n",
      "        WHEN age >= 20 AND age < 30 THEN '20-30'\n",
      "        WHEN age >= 30 AND age < 40 THEN '30-40'\n",
      "        WHEN age >= 40 AND age < 50 THEN '40-50'\n",
      "        WHEN age >= 50 AND age < 60 THEN '50-60'\n",
      "        WHEN age >= 60 AND age < 70 THEN '60-70'\n",
      "        WHEN age >= 70 AND age < 80 THEN '70-80'\n",
      "        ELSE '80+'\n",
      "    END AS age_group,\n",
      "    COUNT(*) AS user_count\n",
      "FROM user\n",
      "GROUP BY age_group;\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT \\n    CASE \\n        WHEN age >= 0 AND age < 10 THEN '0-10'\\n        WHEN age >= 10 AND age < 20 THEN '10-20'\\n        WHEN age >= 20 AND age < 30 THEN '20-30'\\n        WHEN age >= 30 AND age < 40 THEN '30-40'\\n        WHEN age >= 40 AND age < 50 THEN '40-50'\\n        WHEN age >= 50 AND age < 60 THEN '50-60'\\n        WHEN age >= 60 AND age < 70 THEN '60-70'\\n        WHEN age >= 70 AND age < 80 THEN '70-80'\\n        ELSE '80+'\\n    END AS age_group,\\n    COUNT(*) AS user_count\\nFROM user\\nGROUP BY age_group;\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vn.generate_sql('统计一下各年龄段的用户数量,年龄段是指0-10,10-20,20-30,30-40,40-50,50-60,60-70,70-80...左闭右开区间')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding documentation....\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Number of requested results 10 is greater than number of elements in index 2, updating n_results = 2\n",
      "Number of requested results 10 is greater than number of elements in index 1, updating n_results = 1\n",
      "Number of requested results 10 is greater than number of elements in index 3, updating n_results = 3\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SQL Prompt: [{'role': 'system', 'content': 'You are a SQL expert. Please help to generate a SQL query to answer the question. Your response should ONLY be based on the given context and follow the response guidelines and format instructions. \\n===Tables \\nCREATE TABLE IF NOT EXISTS user (\\n        id INT PRIMARY KEY COMMENT \\'用户ID\\' ,\\n        name VARCHAR(100) COMMENT \\'姓名\\',\\n        age INT COMMENT \\'年龄\\'\\n    ) COMMENT \\'用户信息表\\';\\n\\n\\n\\n===Additional Context \\n\\n用户年龄段划分逻辑：0-10,10-20,20-30,30-40,40-50,50-60,60-70,70-80...左闭右开区间\\n\\n用户年龄段划分逻辑：0-10,10-20,20-30,30-40,40-50,50-60,60-70,70-80,80-90,90-100,100以上...左闭右开区间\\n\\n\"福报\"是指age>=35岁，也就是可以向社会输送的人才\\n\\n===Response Guidelines \\n1. If the provided context is sufficient, please generate a valid SQL query without any explanations for the question. \\n2. If the provided context is almost sufficient but requires knowledge of a specific string in a particular column, please generate an intermediate SQL query to find the distinct strings in that column. Prepend the query with a comment saying intermediate_sql \\n3. If the provided context is insufficient, please explain why it can\\'t be generated. \\n4. Please use the most relevant table(s). \\n5. If the question has been asked and answered before, please repeat the answer exactly as it was given before. \\n6. Ensure that the output SQL is SQL-compliant and executable, and free of syntax errors. \\n'}, {'role': 'user', 'content': '小鱼儿的年龄'}, {'role': 'assistant', 'content': 'select age from user where name=\"小鱼儿\"'}, {'role': 'user', 'content': 'Who are the users aged between 10 and 20?'}, {'role': 'assistant', 'content': 'select name from user where age between 10 and 20'}, {'role': 'user', 'content': '统计一下各年龄段的用户数量'}]\n",
      "LLM Response: ```sql\n",
      "SELECT \n",
      "    CASE \n",
      "        WHEN age < 10 THEN '0-10'\n",
      "        WHEN age < 20 THEN '10-20'\n",
      "        WHEN age < 30 THEN '20-30'\n",
      "        WHEN age < 40 THEN '30-40'\n",
      "        WHEN age < 50 THEN '40-50'\n",
      "        WHEN age < 60 THEN '50-60'\n",
      "        WHEN age < 70 THEN '60-70'\n",
      "        WHEN age < 80 THEN '70-80'\n",
      "        WHEN age < 90 THEN '80-90'\n",
      "        WHEN age < 100 THEN '90-100'\n",
      "        ELSE '100以上'\n",
      "    END AS age_group,\n",
      "    COUNT(*) AS user_count\n",
      "FROM user\n",
      "GROUP BY age_group;\n",
      "```\n",
      "Extracted SQL: SELECT \n",
      "    CASE \n",
      "        WHEN age < 10 THEN '0-10'\n",
      "        WHEN age < 20 THEN '10-20'\n",
      "        WHEN age < 30 THEN '20-30'\n",
      "        WHEN age < 40 THEN '30-40'\n",
      "        WHEN age < 50 THEN '40-50'\n",
      "        WHEN age < 60 THEN '50-60'\n",
      "        WHEN age < 70 THEN '60-70'\n",
      "        WHEN age < 80 THEN '70-80'\n",
      "        WHEN age < 90 THEN '80-90'\n",
      "        WHEN age < 100 THEN '90-100'\n",
      "        ELSE '100以上'\n",
      "    END AS age_group,\n",
      "    COUNT(*) AS user_count\n",
      "FROM user\n",
      "GROUP BY age_group;\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT \\n    CASE \\n        WHEN age < 10 THEN '0-10'\\n        WHEN age < 20 THEN '10-20'\\n        WHEN age < 30 THEN '20-30'\\n        WHEN age < 40 THEN '30-40'\\n        WHEN age < 50 THEN '40-50'\\n        WHEN age < 60 THEN '50-60'\\n        WHEN age < 70 THEN '60-70'\\n        WHEN age < 80 THEN '70-80'\\n        WHEN age < 90 THEN '80-90'\\n        WHEN age < 100 THEN '90-100'\\n        ELSE '100以上'\\n    END AS age_group,\\n    COUNT(*) AS user_count\\nFROM user\\nGROUP BY age_group;\""
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "text/plain": [
       "\"SELECT \\n  CASE \\n    WHEN age BETWEEN 0 AND 10 THEN '0-10'\\n    WHEN age BETWEEN 11 AND 20 THEN '11-20'\\n    WHEN age BETWEEN 21 AND 30 THEN '21-30'\\n    WHEN age BETWEEN 31 AND 40 THEN '31-40'\\n    WHEN age BETWEEN 41 AND 50 THEN '41-50'\\n    WHEN age BETWEEN 51 AND 60 THEN '51-60'\\n    WHEN age BETWEEN 61 AND 70 THEN '61-70'\\n    WHEN age BETWEEN 71 AND 80 THEN '71-80'\\n    ELSE '80+'\\n  END AS age_group,\\n  COUNT(*) as user_count\\nFROM user\\nGROUP BY age_group;\""
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 知识沉淀后可以直接提问\n",
    "vn.train(documentation='用户年龄段划分逻辑：0-10,10-20,20-30,30-40,40-50,50-60,60-70,70-80,80-90,90-100,100以上...左闭右开区间')\n",
    "vn.generate_sql('统计一下各年龄段的用户数量')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}